package vip.yuange.gsyzm.utils;

import ai.onnxruntime.*;

import java.awt.image.BufferedImage;
import java.io.IOException;
import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import com.alibaba.fastjson2.JSONArray;

public class YZMOCRUtils {

	public static String ORC(BufferedImage image) {  
		Properties prop = new Properties();  
		try (InputStream input = YZMOCRUtils.class.getClassLoader().getResourceAsStream("application.properties")) {  
			prop.load(input);  
			String onnxModelPath = prop.getProperty("onnx");  
			String charset = prop.getProperty("charset"); 
			OrtEnvironment env = OrtEnvironment.getEnvironment();
			OrtSession session = env.createSession(onnxModelPath);
			JSONArray charsetArray = JSONArray.parseArray(new String(FileUtils.getFileByteArray(charset), "UTF-8"));
			image = PublicUtils.reSize(image, 64 * image.getWidth() / image.getHeight(), 64);
			image = PublicUtils.getGray(image); 
			long[] shape = {1, 1, image.getHeight(), image.getWidth()};
			float[] data = new float[(int)(shape[0] * shape[1] * shape[2] * shape[3])];
			image.getData().getPixels(0, 0, image.getWidth(), image.getHeight(), data);
			for (int i = 0; i < data.length; i++) {
				data[i] /= 255;
				data[i] = (float) ((data[i] - 0.5) / 0.5);
			}
			FloatBuffer floatBuffer = FloatBuffer.wrap(data);  
			OnnxTensor inputTensor = OnnxTensor.createTensor(env, floatBuffer, shape);  
			Map<String, OnnxTensor> inputs = new HashMap<>();  
			inputs.put("input1", inputTensor); 
			OnnxTensor indexTensor = (OnnxTensor) session.run(inputs).get(0);
			long[][] index = (long[][])indexTensor.getValue();
			StringBuilder words = new StringBuilder();
			for (long i: index[0]) {
				words.append((String) charsetArray.get((int) i));
			}
			return words.toString();  
		} catch (IOException | OrtException e) {  
			e.printStackTrace();  
			return null; 
		}  
	}
	
}
